# -*- coding:utf-8 -*-
# @FileName :ddpg.py
# @Time :2024/6/15 下午4:42
# @Author :ShengYe
# @Des :


import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

from my_code.carla_env.carla_env import CarlaEnv
from my_code.rf.rf_utils.replay_buffer import ReplayBuffer
# 定义 Actor 网络
class Actor(nn.Module):
    """
    Actor网络用于生成动作的策略网络。

    Args:
        state_dim (int): 状态空间的维度。
        action_dim (int): 动作空间的维度。
        action_range (float): 动作范围的上下界。

    Attributes:
        fc1, fc2, fc3 (nn.Linear): 神经网络层。
        action_range (float): 动作范围的上下界。

    Methods:
        forward(state): 前向传播函数，生成动作。

    """

    def __init__(self, state_dim, action_dim, action_range):
        super(Actor, self).__init__()
        self.fc1 = nn.Linear(state_dim, 64)
        self.fc2 = nn.Linear(64, 64)
        self.fc3 = nn.Linear(64, action_dim)
        self.action_range = action_range

    def forward(self, state):
        x = F.relu(self.fc1(state))
        x = F.relu(self.fc2(x))
        x = torch.tanh(self.fc3(x)) * self.action_range
        return x


# 定义 Critic 网络
class Critic(nn.Module):
    """
    Critic网络用于评估动作的价值函数网络。

    Args:
        state_dim (int): 状态空间的维度。
        action_dim (int): 动作空间的维度。

    Attributes:
        fc1, fc2, fc3 (nn.Linear): 神经网络层。

    Methods:
        forward(state, action): 前向传播函数，评估动作的价值。
    """

    def __init__(self, state_dim, action_dim):
        super(Critic, self).__init__()
        self.fc1 = nn.Linear(state_dim + action_dim, 64)
        self.fc2 = nn.Linear(64, 64)
        self.fc3 = nn.Linear(64, 1)

    def forward(self, state, action):
        x = torch.cat([state, action], dim=1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


# DDPG Agent
class DDPGAgent:
    """
    DDPG Agent用于训练Actor-Critic模型，并且管理目标网络的软更新。

    Args:
        state_dim (int): 状态空间的维度。
        action_dim (int): 动作空间的维度。
        action_range (float): 动作范围的上下界，默认为1。
        actor_lr (float): Actor网络的学习率。
        critic_lr (float): Critic网络的学习率。
        gamma (float): 折扣因子。
        tau (float): 软更新参数。

    Attributes:
        actor (Actor): 当前策略Actor网络。
        actor_target (Actor): 目标策略Actor网络。
        actor_optimizer (torch.optim): Actor网络的优化器。
        critic (Critic): Critic网络。
        critic_target (Critic): 目标Critic网络。
        critic_optimizer (torch.optim): Critic网络的优化器。
        gamma (float): 折扣因子。
        tau (float): 软更新参数。

    Methods:
        select_action(state): 根据当前状态选择动作。
        train(replay_buffer, batch_size): 使用经验回放训练Actor和Critic网络。
        soft_update(target_net, eval_net, tau): 软更新目标网络的参数。
    """

    def __init__(self, state_dim, action_dim, action_range=1, actor_lr=1e-3, critic_lr=1e-3, gamma=0.99,
                 tau=1e-2):
        self.actor = Actor(state_dim, action_dim, action_range)
        self.actor_target = Actor(state_dim, action_dim, action_range)
        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=actor_lr)

        self.critic = Critic(state_dim, action_dim)
        self.critic_target = Critic(state_dim, action_dim)
        self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=critic_lr)

        self.actor_target.load_state_dict(self.actor.state_dict())
        self.critic_target.load_state_dict(self.critic.state_dict())

        self.gamma = gamma
        self.tau = tau

    def select_action(self, state):
        """
        根据当前状态选择动作。

        Args:
            state (numpy.ndarray): 当前状态数组。

        Returns:
            numpy.ndarray: 选择的动作。
        """
        state = torch.FloatTensor(state).unsqueeze(0)
        return self.actor(state).detach().numpy().flatten()

    def train(self, replay_buffer, batch_size=64):
        """
        使用经验回放训练Actor和Critic网络。

        Args:
            replay_buffer (ReplayBuffer): 经验回放缓冲区。
            batch_size (int): 批处理大小。
        """
        state_batch, action_batch, next_state_batch, reward_batch, done_batch = replay_buffer.sample(batch_size)

        state_batch = torch.FloatTensor(state_batch)
        action_batch = torch.FloatTensor(action_batch)
        next_state_batch = torch.FloatTensor(next_state_batch)
        reward_batch = torch.FloatTensor(reward_batch).unsqueeze(1)
        done_batch = torch.FloatTensor(done_batch).unsqueeze(1)

        # 更新 Critic 网络
        Q = self.critic(state_batch, action_batch)
        next_actions = self.actor_target(next_state_batch)
        Q_next = self.critic_target(next_state_batch, next_actions.detach())
        target_Q = reward_batch + (1 - done_batch) * self.gamma * Q_next
        critic_loss = F.mse_loss(Q, target_Q.detach())

        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()

        # 更新 Actor 网络
        actor_loss = -self.critic(state_batch, self.actor(state_batch)).mean()

        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        self.actor_optimizer.step()

        # 软更新目标网络
        self.soft_update(self.actor_target, self.actor, self.tau)
        self.soft_update(self.critic_target, self.critic, self.tau)

    @staticmethod
    def soft_update(target_net, eval_net, tau):
        """
        软更新目标网络的参数。

        Args:
            target_net (nn.Module): 目标网络。
            eval_net (nn.Module): 当前网络。
            tau (float): 软更新参数。
        """
        for target_param, param in zip(target_net.parameters(), eval_net.parameters()):
            target_param.data.copy_(tau * param.data + (1.0 - tau) * target_param.data)


# 模拟Carla环境

# 经验回放缓冲区


